import os
import winreg
import ctypes
import argparse
from typing import Optional, List, Tuple


class EnvironmentVariableManager:
    """Windows环境变量管理工具"""

    def __init__(self):
        self.scope_options = {
            'user': winreg.HKEY_CURRENT_USER,
            'system': winreg.HKEY_LOCAL_MACHINE
        }
        self.system_env_key_path = r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment"
        self.user_env_key_path = r"Environment"

    def is_admin(self) -> bool:
        """检查脚本是否以管理员权限运行"""
        try:
            return ctypes.windll.shell32.IsUserAnAdmin()
        except:
            return False

    def get_environment_variable(self, name: str, scope: str = "user") -> Optional[str]:
        """
        获取环境变量的值

        Args:
            name: 环境变量名称
            scope: 作用域，"user"（用户）或"system"（系统）

        Returns:
            环境变量的值，如果不存在则返回None
        """
        try:
            key_path = self.system_env_key_path if scope == "system" else self.user_env_key_path
            key = winreg.OpenKey(self.scope_options[scope], key_path)
            value, _ = winreg.QueryValueEx(key, name)
            winreg.CloseKey(key)
            return value
        except WindowsError:
            return None

    def set_environment_variable(self, name: str, value: str, scope: str = "user") -> bool:
        """
        设置环境变量

        Args:
            name: 环境变量名称
            value: 环境变量值
            scope: 作用域，"user"（用户）或"system"（系统，需要管理员权限）

        Returns:
            是否设置成功
        """
        # 检查权限
        if scope == "system" and not self.is_admin():
            print("错误: 设置系统环境变量需要管理员权限")
            return False

        try:
            key_path = self.system_env_key_path if scope == "system" else self.user_env_key_path
            key = winreg.OpenKey(
                self.scope_options[scope],
                key_path,
                0,
                winreg.KEY_SET_VALUE
            )

            # 设置环境变量
            winreg.SetValueEx(key, name, 0, winreg.REG_EXPAND_SZ, value)
            winreg.CloseKey(key)

            # 通知系统环境变量已更改
            self._notify_environment_change()

            print(f"成功设置 {scope} 环境变量: {name}={value}")
            return True
        except Exception as e:
            print(f"设置环境变量失败: {e}")
            return False

    def delete_environment_variable(self, name: str, scope: str = "user") -> bool:
        """
        删除环境变量

        Args:
            name: 环境变量名称
            scope: 作用域，"user"（用户）或"system"（系统，需要管理员权限）

        Returns:
            是否删除成功
        """
        # 检查权限
        if scope == "system" and not self.is_admin():
            print("错误: 删除系统环境变量需要管理员权限")
            return False

        try:
            key_path = self.system_env_key_path if scope == "system" else self.user_env_key_path
            key = winreg.OpenKey(
                self.scope_options[scope],
                key_path,
                0,
                winreg.KEY_SET_VALUE
            )

            # 删除环境变量
            winreg.DeleteValue(key, name)
            winreg.CloseKey(key)

            # 通知系统环境变量已更改
            self._notify_environment_change()

            print(f"成功删除 {scope} 环境变量: {name}")
            return True
        except WindowsError as e:
            if e.winerror == 2:  # 错误码2表示"系统找不到指定的文件"
                print(f"环境变量不存在: {name}")
                return True
            else:
                print(f"删除环境变量失败: {e}")
                return False
        except Exception as e:
            print(f"删除环境变量失败: {e}")
            return False

    def append_to_path(self, new_path: str, scope: str = "user") -> bool:
        """
        向PATH环境变量添加新路径

        Args:
            new_path: 要添加的路径
            scope: 作用域，"user"（用户）或"system"（系统，需要管理员权限）

        Returns:
            是否添加成功
        """
        # 获取当前PATH
        path_value = self.get_environment_variable("PATH", scope) or ""

        # 检查路径是否已存在（不区分大小写）
        path_list = [p.lower() for p in path_value.split(";") if p]
        if new_path.lower() in path_list:
            print(f"路径已存在: {new_path}")
            return True

        # 添加新路径
        new_path_value = path_value + ";" + new_path if path_value else new_path
        return self.set_environment_variable("PATH", new_path_value, scope)

    def remove_from_path(self, path_to_remove: str, scope: str = "user") -> bool:
        """
        从PATH环境变量中移除路径

        Args:
            path_to_remove: 要移除的路径
            scope: 作用域，"user"（用户）或"system"（系统，需要管理员权限）

        Returns:
            是否移除成功
        """
        # 获取当前PATH
        path_value = self.get_environment_variable("PATH", scope) or ""
        path_list = path_value.split(";")

        # 检查路径是否存在（不区分大小写）
        path_to_remove_lower = path_to_remove.lower()
        if not any(p.lower() == path_to_remove_lower for p in path_list if p):
            print(f"路径不存在于PATH中: {path_to_remove}")
            return True

        # 移除路径
        new_path_list = [p for p in path_list if p.lower() != path_to_remove_lower]
        new_path_value = ";".join(new_path_list)

        return self.set_environment_variable("PATH", new_path_value, scope)

    def list_environment_variables(self, scope: str = "user") -> List[Tuple[str, str]]:
        """
        列出所有环境变量

        Args:
            scope: 作用域，"user"（用户）或"system"（系统）

        Returns:
            环境变量列表，格式为[(变量名, 变量值), ...]
        """
        try:
            key_path = self.system_env_key_path if scope == "system" else self.user_env_key_path
            key = winreg.OpenKey(self.scope_options[scope], key_path)

            env_vars = []
            i = 0
            while True:
                try:
                    name, value, _ = winreg.EnumValue(key, i)
                    env_vars.append((name, value))
                    i += 1
                except WindowsError:
                    break

            winreg.CloseKey(key)
            return env_vars
        except Exception as e:
            print(f"列出环境变量失败: {e}")
            return []

    def _notify_environment_change(self) -> None:
        """通知系统环境变量已更改"""
        # 广播WM_SETTINGCHANGE消息
        HWND_BROADCAST = 0xFFFF
        WM_SETTINGCHANGE = 0x001A
        SMTO_ABORTIFHUNG = 0x0002

        ctypes.windll.user32.SendMessageTimeoutW(
            HWND_BROADCAST, WM_SETTINGCHANGE, 0, "Environment",
            SMTO_ABORTIFHUNG, 5000, None
        )


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Windows环境变量管理工具")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # 创建"set"命令
    set_parser = subparsers.add_parser("set", help="设置环境变量")
    set_parser.add_argument("name", help="环境变量名称")
    set_parser.add_argument("value", help="环境变量值")
    set_parser.add_argument("--scope", choices=["user", "system"], default="user", help="作用域，默认为user")

    # 创建"delete"命令
    delete_parser = subparsers.add_parser("delete", help="删除环境变量")
    delete_parser.add_argument("name", help="环境变量名称")
    delete_parser.add_argument("--scope", choices=["user", "system"], default="user", help="作用域，默认为user")

    # 创建"append"命令
    append_parser = subparsers.add_parser("append", help="向PATH添加路径")
    append_parser.add_argument("path", help="要添加的路径")
    append_parser.add_argument("--scope", choices=["user", "system"], default="user", help="作用域，默认为user")

    # 创建"remove"命令
    remove_parser = subparsers.add_parser("remove", help="从PATH移除路径")
    remove_parser.add_argument("path", help="要移除的路径")
    remove_parser.add_argument("--scope", choices=["user", "system"], default="user", help="作用域，默认为user")

    # 创建"list"命令
    list_parser = subparsers.add_parser("list", help="列出环境变量")
    list_parser.add_argument("--scope", choices=["user", "system"], default="user", help="作用域，默认为user")

    # 创建"get"命令
    get_parser = subparsers.add_parser("get", help="获取环境变量的值")
    get_parser.add_argument("name", help="环境变量名称")
    get_parser.add_argument("--scope", choices=["user", "system"], default="user", help="作用域，默认为user")

    args = parser.parse_args()

    manager = EnvironmentVariableManager()

    if args.command == "set":
        manager.set_environment_variable(args.name, args.value, args.scope)
    elif args.command == "delete":
        manager.delete_environment_variable(args.name, args.scope)
    elif args.command == "append":
        manager.append_to_path(args.path, args.scope)
    elif args.command == "remove":
        manager.remove_from_path(args.path, args.scope)
    elif args.command == "list":
        env_vars = manager.list_environment_variables(args.scope)
        print(f"{args.scope} 环境变量:")
        for name, value in env_vars:
            print(f"{name}={value}")
    elif args.command == "get":
        value = manager.get_environment_variable(args.name, args.scope)
        if value is not None:
            print(f"{args.name}={value}")
        else:
            print(f"环境变量不存在: {args.name}")


def set_var_main():
    manager = EnvironmentVariableManager()
    manager.set_environment_variable("ANDROID_HOME", "D:/sdk")
    ret = manager.get_environment_variable("ANDROID_HOME")
    print(ret)

if __name__ == "__main__":
    # main()
    set_var_main()